"""
https://www.cnblogs.com/long5683/p/9692845.html

"""
import cv2 as cv
import numpy as np
import os
import sys
import matplotlib.pyplot as plt
import datetime
from python_ai.common.xcommon import *

np.random.seed(1)
np.set_printoptions(edgeitems=1000)


spr = 3
spc = 5
spn = 0
plt.figure(figsize=[14, 7])


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


sep('Load')
dir = '../../../../../large_data/pic/watershed/'
name = 'poker_cards.png'
# dir = '../../../../../large_data/pic/'
# name = 'dog_bird.png'
path = os.path.join(dir, name)
src = cv.imread(path, cv.IMREAD_COLOR)
src_ = src.copy()

my_show_img(src, 'input image', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
cv.imshow('1 input', src)

# 白色背景变成黑色
sep('白色背景变成黑色')
src_gray = cv.cvtColor(src, cv.COLOR_BGR2GRAY)
idx = src_gray == 255
src[idx] = (0, 0, 0)
my_show_img(src, 'black background', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
cv.imshow('2 255->0', src)

# sharpen(提高对比度)
sep('sharpen(提高对比度)')
kernel = np.float32([
    [0, 1, 0],
    [1, -8, 1],
    [0, 1, 0]
])
# 拉普拉斯算子实现边缘提取
imgLaplance = cv.filter2D(src, cv.CV_32F, kernel, borderType=cv.BORDER_DEFAULT)  # 拉普拉斯有浮点数计算，位数要提高到32
# check and draw imgLaplance
print_numpy_ndarray_info(imgLaplance, 'imgLaplance')
imgLaplance_ = imgLaplance.copy()
imgLaplance_ = cv.normalize(imgLaplance_, None, 0., 255., norm_type=cv.NORM_MINMAX)
imgLaplance_ = imgLaplance_.astype(np.uint8)
my_show_img(imgLaplance_, 'imgLaplance', cmap='gray')
cv.imshow('3 Laplance', imgLaplance_)

# 原图减边缘（白色）实现边缘增强
sep('原图减边缘（白色）实现边缘增强')
sharpenImg = src.astype(np.float32)
resultImg = cv.subtract(sharpenImg, imgLaplance)
# check and draw resultImg
print_numpy_ndarray_info(resultImg, 'resultImg')
resultImg_ = resultImg.copy()
cv.normalize(resultImg_, resultImg_, 0., 255., norm_type=cv.NORM_MINMAX)
resultImg_ = resultImg_.astype(np.uint8)
my_show_img(resultImg_, 'subtract res', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
cv.imshow('4 subtract', resultImg_)

# 转换成二值图
sep('转换成二值图')
resultImg = cv.cvtColor(resultImg_, cv.COLOR_BGR2GRAY)
ret, binary = cv.threshold(resultImg, 40, 255, cv.THRESH_OTSU + cv.THRESH_BINARY)
my_show_img(binary, 'binary', cmap='gray')
cv.imshow('5 binary', binary)

# 距离变换
sep('距离变换')
distImg = cv.distanceTransform(binary, cv.DIST_L2, 3, dstType=cv.CV_32F)
print_numpy_ndarray_info(distImg, 'distImg after distance trans')
cv.normalize(distImg, distImg, 0., 1., cv.NORM_MINMAX)
print_numpy_ndarray_info(distImg, 'distImg after normalization')
my_show_img(distImg, 'distance', cmap='gray')
cv.imshow('6 distance', distImg)

# 二值化
sep('二值化')
ret, distImg = cv.threshold(distImg, 0.4, 1.0, cv.THRESH_BINARY)
print_numpy_ndarray_info(distImg, 'distImg after threshold')
my_show_img(distImg, 'bin', cmap='gray')
cv.imshow('7 bin', distImg)

# 腐蚀(使得连在一起的部分分开)
sep('腐蚀(使得连在一起的部分分开)')
k1 = np.ones([3, 3], dtype=np.uint8)
distImg = cv.erode(distImg, k1)
print_numpy_ndarray_info(distImg, 'distImg after erosion')
my_show_img(distImg, 'erode', cmap='gray')
cv.imshow('7 erode', distImg)

# 标记
sep('标记')
distImg *= 255.
dist_8u = distImg.astype(np.uint8)
print_numpy_ndarray_info(dist_8u, 'dist_8u')
contours, hierarchy = cv.findContours(dist_8u, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE, offset=(0, 0))
bg = np.zeros(src.shape, dtype=np.uint8)
cv.drawContours(bg, contours, -1, (0, 255, 0), 1)
my_show_img(bg, 'contours', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
cv.imshow('8 contours', bg)
# 创建标记
marker = np.zeros(distImg.shape, dtype=np.int32)
# 画标记
for i, c in enumerate(contours):
    cv.drawContours(marker, [c], 0, i + 1, -1)
# 可视化标记
# color table start
# 随机生成几种颜色
colorTab = [(255, 255, 255)]
for i in range(len(contours)):
    b = np.random.randint(0, 256)
    g = np.random.randint(0, 256)
    r = np.random.randint(0, 256)
    colorTab.append((b, g, r))
colorTab = np.uint8(colorTab)
# color table end
print(np.unique(marker))
marker_drawing = colorTab[marker]
my_show_img(marker_drawing, 'marker', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
cv.imshow('8 marker', marker_drawing)

# 分水岭变换
sep('分水岭变换')
src = src_.copy()
cv.watershed(src, marker)
print(np.unique(marker))
marker_drawing = marker
marker_drawing[marker_drawing == -1] = 0
print(np.unique(marker_drawing))
marker_drawing = colorTab[marker]
my_show_img(marker_drawing, 'after water', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
cv.imshow('9 after water', marker_drawing)

plt.show()
cv.waitKey(0)
cv.destroyAllWindows()
